from threading import Condition
from numpy._core.multiarray import scalar
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report
from joblib import dump

#读取数据
df = pd.read_csv('daily.csv')

#计算收益和风险比率
df['max_ratio'] = df['max_close'] / df['the_close']
df['min_ratio'] = df['min_close'] / df['the_close']

#自动划分高低阀值
high_retuen_threshold = df['max_ratio'].quantile(0.4)
high_risk_threshold = df['min_ratio'].quantile(0.4)

#生成分类标签
conditions = [
    (df['max_ratio'] >= high_retuen_threshold) & (df['min_ratio'] <= high_risk_threshold),
    (df['max_ratio'] <= high_retuen_threshold) & (df['min_ratio'] > high_risk_threshold),
    (df['max_ratio'] > high_retuen_threshold) & (df['min_ratio'] > high_risk_threshold),
    (df['max_ratio'] < high_retuen_threshold) & (df['min_ratio'] <= high_risk_threshold),
]
labels = ['高收益高风险','高收益低风险','低收益低风险','低收益高风险']
df['category'] = np.select(conditions, labels, default='未知')

#特征选择
features = df[[
    'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
    'fcff', 'fcfe','tangible_asset', 'bps', 'grossprofit_margin', 'npta', 'roic'
]]
#标签编码
from sklearn.preprocessing import LabelEncoder
le = LabelEncoder()
df['category_encoded'] = le.fit_transform(df['category'])

#数据标准化
scaler = StandardScaler()
X = scalar.fit_transform(features)
y = df['category_encoded']

#划分训练集合测试集
X_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

#训练KNN模型
knn = KNeighborsClassifier(n_neighbors=2)
knn.fit(X_train, y_train)

#预估与评估
y_pred = knn.predict(x_test)
print(classification_report(y_test, y_pred, target_names=le.classes_))

#加载模型和预处理对象
dump(knn, 'knn_classifier.joblib')
dump(scaler, 'feature_scaler.joblib')
dump(le, 'label_encoder.joblib')